import numpy as np
class KMeans:
    def __init__(self,data,K):
        self.data = data
        self.K = K
        self.num_samples = self.data.shape[0]
        self.num_features = self.data.shape[1]


    def train(self,num_iterations):
        self.centroids_init()
        closet_centroids_idx = np.zeros((self.num_samples,1))
        for i in range(num_iterations):
            closet_centroids_idx = self.select_clusters()
            self.calculate_centroids(closet_centroids_idx)
        return self.centroids,closet_centroids_idx

    def centroids_init(self):
        random_idxs = np.random.permutation(self.num_samples)
        self.centroids = self.data[random_idxs[:self.K],:]
    
    def select_clusters(self):
        sample_clusters = np.zeros((self.num_samples,1))
        for sample_idx,sample in enumerate(self.data):
            distances = np.zeros((self.K,1))
            for centroid_idx,centroid in enumerate(self.centroids):
                distance = np.dot((sample-centroid).T,sample-centroid)
                distances[centroid_idx] = distance
            sample_clusters[sample_idx] = np.argmin(distances) 
        return sample_clusters
    
    def calculate_centroids(self,closet_centroids_idx):
        for k in range(self.K):
            closet_idx = closet_centroids_idx==k
            samples = self.data[closet_idx.flatten(),:]
            self.centroids[k] = np.mean(samples,axis=0)
